package dlock

import (
	"context"
	"fmt"
	"github.com/go-redis/redis/v8"
	"github.com/google/uuid"
	"strconv"
	"sync"
	"time"
)

// 可重入锁的Lua脚本
const reentrantLockScript = `
local key = KEYS[1]
local threadId = ARGV[1]
local ttl = tonumber(ARGV[2])

-- 如果锁不存在
if redis.call('EXISTS', key) == 0 then
    -- 创建锁，初始计数为1
    redis.call('HSET', key, threadId, 1)
    redis.call('EXPIRE', key, ttl)
    return 1
end

-- 如果锁存在且是当前线程持有
if redis.call('HEXISTS', key, threadId) == 1 then
    -- 增加计数
    redis.call('HINCRBY', key, threadId, 1)
    -- 刷新过期时间
    redis.call('EXPIRE', key, ttl)
    return 1
end

-- 锁被其他线程持有
return 0
`

// 释放可重入锁的Lua脚本
const releaseReentrantLockScript = `
local key = KEYS[1]
local threadId = ARGV[1]

-- 如果锁不存在
if redis.call('EXISTS', key) == 0 then
    return 0
end

-- 如果锁不是当前线程持有
if redis.call('HEXISTS', key, threadId) == 0 then
    return 0
end

-- 减少计数
local count = redis.call('HINCRBY', key, threadId, -1)

-- 如果计数为0，删除锁
if count == 0 then
    redis.call('DEL', key)
    return 1
else
    return 1
end
`

// ReentrantLock 可重入分布式锁
type ReentrantLock struct {
	client    *redis.Client
	prefix    string
	threadID  string
	lockKeys  map[string]bool
	keysMutex sync.Mutex
}

// NewReentrantLock 创建可重入分布式锁
func NewReentrantLock(client *redis.Client, prefix string) *ReentrantLock {
	if prefix == "" {
		prefix = "dlock:reentrant:"
	}
	
	return &ReentrantLock{
		client:   client,
		prefix:   prefix,
		threadID: uuid.New().String(),
		lockKeys: make(map[string]bool),
	}
}

// Lock 获取锁，支持可重入
func (rl *ReentrantLock) Lock(ctx context.Context, key string, ttl time.Duration) (string, error) {
	lockKey := rl.prefix + key
	
	// 尝试获取锁
	retryInterval := 100 * time.Millisecond
	timeout := 5 * time.Second
	
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	
	for {
		// 执行Lua脚本
		result, err := rl.client.Eval(ctx, reentrantLockScript, []string{lockKey}, rl.threadID, int(ttl.Seconds())).Result()
		if err != nil {
			return "", err
		}
		
		// 获取锁成功
		if result.(int64) == 1 {
			// 记录已获取的锁
			rl.keysMutex.Lock()
			rl.lockKeys[key] = true
			rl.keysMutex.Unlock()
			
			return rl.threadID, nil
		}
		
		// 等待一段时间后重试
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-timer.C:
			return "", ErrLockAcquireFailed
		case <-time.After(retryInterval):
			// 继续重试
		}
	}
}

// TryLock 尝试获取锁，立即返回结果，支持可重入
func (rl *ReentrantLock) TryLock(ctx context.Context, key string, ttl time.Duration) (string, error) {
	lockKey := rl.prefix + key
	
	// 执行Lua脚本
	result, err := rl.client.Eval(ctx, reentrantLockScript, []string{lockKey}, rl.threadID, int(ttl.Seconds())).Result()
	if err != nil {
		return "", err
	}
	
	// 获取锁成功
	if result.(int64) == 1 {
		// 记录已获取的锁
		rl.keysMutex.Lock()
		rl.lockKeys[key] = true
		rl.keysMutex.Unlock()
		
		return rl.threadID, nil
	}
	
	return "", ErrLockAcquireFailed
}

// Unlock 释放锁，支持可重入
func (rl *ReentrantLock) Unlock(ctx context.Context, key string, value string) error {
	// 检查是否是当前线程的锁
	if value != rl.threadID {
		return ErrLockReleaseFailed
	}
	
	lockKey := rl.prefix + key
	
	// 执行Lua脚本
	result, err := rl.client.Eval(ctx, releaseReentrantLockScript, []string{lockKey}, rl.threadID).Result()
	if err != nil {
		return err
	}
	
	if result.(int64) != 1 {
		return ErrLockReleaseFailed
	}
	
	// 移除已释放的锁记录
	rl.keysMutex.Lock()
	delete(rl.lockKeys, key)
	rl.keysMutex.Unlock()
	
	return nil
}

// GetLockCount 获取当前锁的重入次数
func (rl *ReentrantLock) GetLockCount(ctx context.Context, key string) (int, error) {
	lockKey := rl.prefix + key
	
	// 检查锁是否存在
	exists, err := rl.client.Exists(ctx, lockKey).Result()
	if err != nil {
		return 0, err
	}
	
	if exists == 0 {
		return 0, nil
	}
	
	// 获取当前线程的计数
	count, err := rl.client.HGet(ctx, lockKey, rl.threadID).Result()
	if err != nil {
		if err == redis.Nil {
			return 0, nil // 当前线程没有持有锁
		}
		return 0, err
	}
	
	// 转换为整数
	countInt, err := strconv.Atoi(count)
	if err != nil {
		return 0, err
	}
	
	return countInt, nil
}

// UnlockAll 释放当前线程持有的所有锁
func (rl *ReentrantLock) UnlockAll(ctx context.Context) error {
	rl.keysMutex.Lock()
	keys := make([]string, 0, len(rl.lockKeys))
	for key := range rl.lockKeys {
		keys = append(keys, key)
	}
	rl.keysMutex.Unlock()
	
	for _, key := range keys {
		err := rl.Unlock(ctx, key, rl.threadID)
		if err != nil {
			return fmt.Errorf("释放锁[%s]失败: %w", key, err)
		}
	}
	
	return nil
}